#include "shp/Shp.h"
#include <fstream>
#include <list>
#include "common/exception.h"
#include "WW_RLE_Zero.h"

using namespace std;

class streamHelper
{
private:
    list<vector<uint8_t>> buffer;
    size_t size;
public:
    streamHelper()
    {
        buffer = list<vector<uint8_t>>();
        size = 0;
    }
    [[nodiscard]] size_t Size() const { return size; }
    void operator<<(vector<uint8_t>& c)
    {
        size += c.size();
        buffer.emplace_back(vector<uint8_t>());
        auto iter = buffer.end();
        iter--;
        iter->swap(c);
    }
    void operator<<(uint16_t n)
    {
        size += 2;
        auto p = (uint8_t*)&n;
        buffer.emplace_back(vector<uint8_t>({*(p), *(p + 1)}));
    }
    void operator>>(ostream& s)
    {
        for (auto i : buffer)
            s.write((char*)i.data(), (int64_t)i.size());
    }
};

Shp::Shp(const Shp::ShpHeader &header_)
{
    header = header_;
    frames = vector<ShpFrame>();
    frames.reserve(header_.frameTotal);
}

shared_ptr<Shp> Shp::Load(const std::string_view &path)
{
    ifstream ifs;
    ifs.open(path.data(), ios_base::binary);
    if (!ifs.is_open())
        throw BadException(string("Could open file: ") + path.data());

    ShpHeader header_{};
    ifs.read((char*)(&(header_.empty)), 2);
    if (header_.empty != 0)
        throw BadException("This file is not a TSShp format.");
    ifs.read((char*)(&(header_.width)), 2);
    ifs.read((char*)(&(header_.height)), 2);
    ifs.read((char*)(&(header_.frameTotal)), 2);

    auto newShp = make_shared<Shp>(header_);

    for (int i = 0; i < header_.frameTotal; ++i)
    {
        newShp->frames.emplace_back(ShpFrame{});
        ShpFrame& frame = newShp->frames[i];
        ifs.read((char*)(&(frame.x)), 2);
        ifs.read((char*)(&(frame.y)), 2);
        ifs.read((char*)(&(frame.width)), 2);
        ifs.read((char*)(&(frame.height)), 2);
        ifs.read((char*)(&(frame.flags)), 4);
        ifs.read((char*)(frame.frameColor), 4);
        ifs.read((char*)(&(frame.reserved)), 4);
        ifs.read((char*)(&(frame.dataOffset)), 4);
        auto pos = ifs.tellg();
        frame.data = make_shared<Uint8Seq>(frame.width * frame.height);
        ifs.seekg(frame.dataOffset);
        if(frame.flags == 3)
        {
            for (int k = 0; k < frame.height; ++k)
            {
                uint16_t length;
                ifs.read((char*)(&length), 2);
                length -= 2;
                auto src = new uint8_t[length];
                ifs.read((char*)src, length);
                WW_RLE_Zero::Decode(src, frame.data->ptr() + k * frame.width, length, frame.width);
                delete[] src;
            }
        }
        else
        {
            ifs.read((char*)(frame.data->ptr()), frame.width * frame.height);
        }
        ifs.seekg(pos);
    }
    ifs.close();
    return newShp;
}

void Shp::Write(const string_view &path)
{
    ofstream ofs;
    ofs.open(path.data(), ios_base::binary | ios_base::trunc);
    if (!ofs.is_open())
        throw BadException(string("Could open file: ") + path.data());
    ofs.write((char*)&header.empty, 2);
    ofs.write((char*)&header.width, 2);
    ofs.write((char*)&header.height, 2);
    ofs.write((char*)&header.frameTotal, 2);
    uint32_t offsetBegin = 8 + 24 * header.frameTotal;
    vector<streamHelper> dataC;
    dataC.reserve(header.frameTotal);
    for (int i = 0; i < header.frameTotal; ++i)
    {
        ShpFrame& frame = frames[i];
        if(frame.flags == 3)
        {
            dataC.emplace_back(streamHelper());
            for (int k = 0; k < frame.height; ++k)
            {
                vector<uint8_t> out;
                WW_RLE_Zero::Encode(frame.data->ptr() + frame.width * k, frame.width, out);
                dataC[i] << (uint16_t) out.size() + 2;
                dataC[i] << out;
            }
        }
        if (i == 0)
            frame.dataOffset = offsetBegin;
        else if(frame.flags == 3)
            frame.dataOffset = frames[i - 1].dataOffset + dataC[i - 1].Size();
        else
            frame.dataOffset = frames[i - 1].dataOffset + frames[i - 1].width * frames[i - 1].height;
        ofs.write((char*)&frame.x, 2);
        ofs.write((char*)&frame.y, 2);
        ofs.write((char*)&frame.width, 2);
        ofs.write((char*)&frame.height, 2);
        ofs.write((char*)&frame.flags, 4);
        ofs.write((char*)&frame.frameColor, 4);
        ofs.write((char*)&frame.reserved, 4);
        ofs.write((char*)&frame.dataOffset, 4);
    }
    for (int i = 0; i < header.frameTotal; ++i)
    {
        if (frames[i].flags == 3)
            dataC[i] >> ofs;
        else
            ofs.write((char*)(frames[i].data->ptr()), frames[i].width * frames[i].height);
    }
    ofs.close();
}


